import warnings
import pandas as pd
import joblib


def health_data_conversion(data):
    """
    标准化数据
    :param data: data
    :return: 标准化后的数据
    """
    # 给定的列顺序
    warnings.filterwarnings('ignore')
    pd.set_option('display.max_columns', None)
    health_columns = joblib.load('./model/health_model/health_columns.pkl')
    print(health_columns)
    print("===================================")
    # 创建空的 DataFrame
    new_data = pd.DataFrame(columns=health_columns)

    print(new_data.columns)

    # 将原始数据转换为新的 DataFrame
    for key, value in vars(data).items():
        if key == 'gender' or key == 'ever_married' or key == 'work_type' or key == 'residence_type' or\
                key == 'smoking_status':
            new_data[value] = [1]
        else:
            if key != 'mode':
                new_data[key] = value
    new_data.fillna(0, inplace=True)

    print(new_data)

    return new_data


def salary_data_conversion(data):
    """
    标准化数据
    :param data: data
    :return: 标准化后的数据
    """
    # 给定的列顺序
    warnings.filterwarnings('ignore')
    pd.set_option('display.max_columns', None)
    salary_columns = joblib.load('./model/salary_model/salary_columns.pkl')
    print(salary_columns)
    print("===================================")
    # 创建空的 DataFrame
    new_data = pd.DataFrame(columns=salary_columns)

    print(new_data.columns)

    # 将原始数据转换为新的 DataFrame
    for key, value in vars(data).items():
        if key == 'age' or key == 'fnlwgt' or key == 'education_num' or key == 'capital_gain' or\
                key == 'capital_loss' or key == 'hours_per_week':
            new_data[key.replace("_", "-")] = value
        else:
            if key != 'mode':
                new_data[key.replace("_", "-") + "_ " + value] = [1]
    new_data.fillna(0, inplace=True)

    print(new_data)

    return new_data
